

import pandas as pd
import os
import numpy as np

import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
import tqdm
from global_land_mask import globe
from gsj_typhoon import tydat,see,count_rapidgrow,split_str_id,tydat_NH
from geopy.distance import geodesic




names,track_ids = ['2023NH42_36','2023NH47_46'],[42,47]
tynames,tyids = split_str_id(names)


for trackid,ty,tyid in zip(track_ids,tynames,tyids):
    RIstd = 7
    dt = timedelta(days=2)
    directory =  f'D:\\met_data\\ty_ensemble\\{ty}_{tyid}\\'    # 要筛选的文件开头
    dates_name = os.listdir(directory)
    dates_paths = [directory + date for date in dates_name]
    obs_path = r'D:\met_data\ty_obs\storms\storms_NH2023'
    pic_savepath = rf'D:\met_data\pics\{ty}_{tyid}-RI{RIstd}\RIpoints'
    os.makedirs(pic_savepath,exist_ok=True)
    
    # 有底图那种
    def track(ty,ax,lon,lat,title=None,color=None,marker=".",label=None,show_label=False,s=12):
        ax.add_feature(cfeature.COASTLINE,linewidth=0.5) #海岸线
        ax.add_feature(cfeature.BORDERS, linestyle=':')   #国界
        china_provinces = cfeature.NaturalEarthFeature(category='cultural',name='admin_1_states_provinces_lines', scale='10m', facecolor='none')
        ax.add_feature(china_provinces, edgecolor='black', linewidth=0.5) #省界
        ax.set(xlim=(ty.extent[0],ty.extent[1]))
        ax.set(ylim=(ty.extent[2],ty.extent[3]))    
        #这种设定方法适用于这样子给出范围。
        ax.set_xticks(np.linspace(ty.extent[0], ty.extent[1]+5, 5))
        ax.set_yticks(np.linspace(ty.extent[2], ty.extent[3]+5, 5))
        ax.xaxis.set_major_formatter(LongitudeFormatter(number_format='.1f'))
        ax.xaxis.set_minor_locator(plt.MultipleLocator(1))
        ax.yaxis.set_major_formatter(LatitudeFormatter(number_format='.1f'))
        ax.yaxis.set_minor_locator(plt.MultipleLocator(1))
        ax.tick_params(axis='both', labelsize=6, direction='out')
        ###
        now = start
        i=0
        colors = plt.cm.tab20(np.linspace(0, 1, 19))
        while now<=ty.time[-1]:
            end = now+dt
            want = (ty.time>=now) & (ty.time<end) & (record==1)
            label_to_use = f"{dt.days*i}-{dt.days*(i+1)}d" if show_label else None
            ax.scatter(lon[want],lat[want],marker=marker,transform=ccrs.PlateCarree(),color=colors[i],label=label_to_use,s=s)
            now=end ; i+=1
    
    
    def track_of_members(ax ,ty,time,lat,lon,marker='.',show_label=True,s=24):
        # map
        ax.add_feature(cfeature.COASTLINE,linewidth=0.5) #海岸线
        ax.add_feature(cfeature.BORDERS, linestyle=':')   #国界
        china_provinces = cfeature.NaturalEarthFeature(category='cultural',name='admin_1_states_provinces_lines', scale='10m', facecolor='none')
        ax.add_feature(china_provinces, edgecolor='black', linewidth=0.5) #省界
        ax.set(xlim=(ty.extent[0],ty.extent[1]))
        ax.set(ylim=(ty.extent[2],ty.extent[3]))    
        ax.set_xticks(np.linspace(ty.extent[0], ty.extent[1]+5, 5))
        ax.set_yticks(np.linspace(ty.extent[2], ty.extent[3]+5, 5))
        ax.xaxis.set_major_formatter(LongitudeFormatter(number_format='.1f'))
        ax.xaxis.set_minor_locator(plt.MultipleLocator(1))
        ax.yaxis.set_major_formatter(LatitudeFormatter(number_format='.1f'))
        ax.yaxis.set_minor_locator(plt.MultipleLocator(1))
        ax.tick_params(axis='both', labelsize=6, direction='out')
        
        # scatter
        now = start
        i=0
        colors = plt.cm.tab20(np.linspace(0, 1, 19))
        time = np.array(time)
        lat = np.array(lat)
        lon = np.array(lon)
        while now<=start+timedelta(days=16):
            end = now+dt
            want = (time>=now) & (time<end)
            label_to_use = f"{dt.days*i}-{dt.days*(i+1)}d" if show_label else None
            ax.scatter(lon[want],lat[want],marker=marker,transform=ccrs.PlateCarree(),color=colors[i],label=label_to_use,s=s)
            now=end ; i+=1
    
    # 起报时间循环
    for i in tqdm.tqdm(range(len(dates_paths)), desc=f"Processing {ty}"):
        try:
            members_paths = [os.path.join(dates_paths[i], f) for f in os.listdir(dates_paths[i]) if f != "TRACK_ID_0" and f.startswith("TRACK_ID")]   # 打印所有符合后缀的文件名
            era_path = [os.path.join(dates_paths[i], f) for f in os.listdir(dates_paths[i]) if f.startswith("TRACK_ID_0")][0]  # 打印所有符合后缀的文件名
            fig, ax = plt.subplots(1, 1, figsize=(10, 10), subplot_kw={'projection': ccrs.PlateCarree()})
            start = datetime.strptime(dates_name[i],"%Y%m%d%H")
            '''draw members'''
            # 循环 某一起报时间的成员
            time=[]
            lat =[]
            lon =[]
            for j in range(len(members_paths)):
                member = tydat(members_paths[j],RIstd)
                record = member.num_rapidgrow()
                RI = (record==1)
                if RI.any():
                    time.extend(member.time[RI].tolist())
                    lat.extend(member.lat[RI].tolist())
                    lon.extend(member.lon[RI].tolist())
                    
            
            track_of_members(ax ,member,time, lat, lon,marker="o",show_label=True)
            
            ''' draw era5 '''
            era = tydat(era_path,RIstd)
            record = era.num_rapidgrow()
            track(era, ax, era.lon, era.lat,marker="s",s=48)
            
            ''' obs '''
            obs = tydat_NH(obs_path,trackid)
            record = count_rapidgrow(RIstd,obs.umax, obs.time)
            track(obs, ax, obs.lon, obs.lat,marker="^",s=72)
            
            # obs = tydat_JMA(obs_path)
            # record = count_rapidgrow(obs.umax, obs.time)
            # if np.sum(record) != 0.0:
            #     lon = obs.lon[record == 1]
            #     lat = obs.lat[record == 1]
            #     track(obs, ax, lon, lat,color="blue",label="obs")
            
            title = dates_name[i]+f"{ty}"
            ax.set_title(title)
            ax.legend()
            
            ax.text(x=0.8,y=0.5,s="square for era5", transform=ax.transAxes,color="blue",)
            ax.text(x=0.8,y=0.45,s="triangle for obs", transform=ax.transAxes,color="blue")
            # plt.savefig(f"D:\\met_data\\pics\\{dates_name[i]}.png") ; plt.close()
        
            
            
            ''' connect obs and member in the same time '''
            # obs = tydat_CMA(obs_path)
            # record_obs = count_rapidgrow(RIstd, obs.umax, obs.time)
            # RItime_obs,RIlat_obs,RIlon_obs = obs.time[record_obs==1],obs.lat[record_obs==1],obs.lon[record_obs==1]
            
            # def neartime(t,times):
            #     '''
            #     t是一个datetime,times是一个datetime数组
            #     返回t在times中距离最近的datetime
            #     '''
            #     closest_time = min(times, key=lambda x: abs(t - x))
            #     return closest_time
                
            # for t in range(len(RItime_obs)):
            #     obstime = RItime_obs[t]
            #     for j in range(len(members_paths)):
            #         member = tydat(members_paths[j],RIstd)
            #         record_member = member.num_rapidgrow()
            #         want = (record_member==1) & (member.time==neartime(obstime, member.time))
            #         if want.any() == False:
            #             continue
            #         else:
            #             print(i,t,j,'yes')
            #             lat_member = member.lat[want][0]
            #             lon_member = member.lon[want][0]
            #             lat_obs    = RIlat_obs[t]
            #             lon_obs    = RIlon_obs[t]
            #             dis = geodesic( (lat_member,lon_member),(lat_obs,lon_obs) ).kilometers
            #             ax.plot( [lon_obs,lon_member],[lat_obs,lat_member],color='lime')
            #             ax.text((lon_member + lon_obs)/2, (lat_member + lat_obs)/2, f'{int(dis)}km', fontsize=5,color="black")
            
            plt.savefig(f'{pic_savepath}\\{title}')
            # plt.show()
            plt.close()
        except Exception as e:
            print(f'{e}')
            




